#include "go_generator.h"
#include "pp_utils.h"
#include "pp_generator.h"

#include <iostream>
#include <sstream>
#include <map>
#include <algorithm>

#define PP_GEN_GO_VERSION "1.0.0"

namespace pp
{
    const std::map<std::string, std::string> TYPES_MAP = {
        {"uint8", "uint8"}, {"int8", "int8"},
        {"uint16", "uint16"}, {"int16", "int16"},
        {"uint32", "uint32"}, {"int32", "int32"},
        {"uint64", "uint64"}, {"int64", "int64"},
        {"float32", "float32"}, {"float64", "float64"},
        {"float", "float32"}, {"double", "float64"},
        {"string", "string"},
    };

    const std::map<std::string, std::string> UNMARSHAL_MAP = {
        {"uint8", "Uint8FieldUnmarshal"}, {"int8", "Int8FieldUnmarshal"},
        {"uint16", "Uint16FieldUnmarshal"}, {"int16", "Int16FieldUnmarshal"},
        {"uint32", "Uint32FieldUnmarshal"}, {"int32", "Int32FieldUnmarshal"},
        {"uint64", "Uint64FieldUnmarshal"}, {"int64", "Int64FieldUnmarshal"},
        {"float32", "Float32FieldUnmarshal"}, {"float64", "Float64FieldUnmarshal"},
        {"float", "Float32FieldUnmarshal"}, {"double", "Float64FieldUnmarshal"},
    };

    std::string GoGenerator::GetGoType(const std::string& ppType)
    {
        auto iter = TYPES_MAP.find(ppType);
        return iter != TYPES_MAP.end() ? iter->second : ppType;
    }

    void GoGenerator::Reset(const std::string& srcfile)
    {
        CodeGenerator::Reset(srcfile);
        m_sGoPackage = "";
    }

    // 处理pp文件中的import定义
    void GoGenerator::HandleImport(const std::string& imports)
    {
        std::string srcfile = m_sSrcDir + "/" + imports;
        std::shared_ptr<GoGenerator> goGen = std::make_shared<GoGenerator>(m_sSrcDir, m_sOutDir);
        std::shared_ptr<CodeGenerator> gen = goGen;
        std::shared_ptr<CodeHandler> hdr = std::make_shared<CodeHandler>(gen);
        Parser parser;
        goGen->Reset(srcfile);
        parser.SetHandler(hdr);
        parser.ParseFile(srcfile);
        m_oImports.push_back(goGen->GetGoPackage());

        // 枚举
        for (auto& v : goGen->GetEnums())
        {
            m_oEnumList.insert(goGen->GetPackage() + "." + v->GetName());
        }
        // 结构体
        for (auto& v : goGen->GetStructs())
        {
            m_oStructMap[goGen->GetPackage() + "." + v->GetName()] = goGen->GetPackage() + ".New" + v->GetName();
        }
    }

    // 处理pp文件中的option定义，忽略不支持的option。比如option go_package是一个option
    void GoGenerator::HandleOption(const std::string& name, const std::string& value)
    {
        if (name == "go_package")
        {
            m_sGoPackage = value;
        }
    }

    // 处理pp文件中的常量定义
    void GoGenerator::HandleConst(const std::shared_ptr<FieldInfo>& fieldInfo)
    {
        CodeGenerator::HandleConst(fieldInfo);
    }

    // 处理pp文件中的enum定义
    void GoGenerator::HandleEnum(const std::shared_ptr<StructInfo>& structInfo)
    {
        CodeGenerator::HandleEnum(structInfo);

        // 枚举缓存
        m_oEnumList.insert(structInfo->GetName());
    }

    // 处理pp文件中的message定义
    void GoGenerator::HandleStruct(const std::shared_ptr<StructInfo>& structInfo)
    {
        CodeGenerator::HandleStruct(structInfo);

        // 结构体缓存
        m_oStructMap[structInfo->GetName()] = "New" + structInfo->GetName();
    }

    // 生成文件
    void GoGenerator::Generate()
    {
        std::string filename = GetFileName(m_sSrcFile);
        std::string relpath = GetRelativePath(m_sSrcDir, m_sSrcFile) + GetGoPackage();
        std::string outpath = m_sOutDir + "/" + relpath + "/";

        try
        {
            CreateDirectories(outpath);
            // 生成go文件
            std::cout << "    --> " << outpath << filename << ".go" << std::endl;
            GenerateGoFile(outpath, filename);
        }
        catch (const std::exception& e)
        {
            std::cerr << M_ERR_COLOR_START "error: source code for file '" << m_sSrcFile << "' fail:"
                << e.what() << M_ERR_COLOR_END << std::endl;
        }
    }

    // 生成go文件内容
    void GoGenerator::GenerateGoFile(const std::string& outdir, const std::string& filename)
    {
        std::string cppfile = outdir + filename + ".go";

        // 打开文件
        std::ofstream ofs(cppfile, std::ios::out | std::ios::trunc);

        if (!ofs.is_open())
        {
            std::cerr << M_ERR_COLOR_START << "error: fail to open file '" << cppfile
                << "' to write" << M_ERR_COLOR_END << std::endl;
            return;
        }

        ofs << "// Code generated by pp_gen_go. DO NOT EDIT." << std::endl
            << "// versions:" << std::endl
            << "//     pp_gen_go " PP_GEN_GO_VERSION << std::endl
            << "//     pp_gen     " PP_GEN_VERSION << std::endl
            << "// source: " << TrimPath(m_sSrcFile) << std::endl << std::endl
            << "package " << m_sPackage << std::endl << std::endl;

        GenerateImport(ofs);    // import
        GenerateConst(ofs);     // const
        GenerateEnum(ofs);      // enum
        GenerateStruct(ofs);    // struct
    }

    // 生成import
    void GoGenerator::GenerateImport(std::ofstream& ofs)
    {
        if (m_oConsts.empty() && m_oEnums.empty() && m_oStructs.empty()) return;

        ofs << "import (" << std::endl;
        ofs << "\t\"bytes\"" << std::endl;

        // 排序
        std::sort(m_oImports.begin(), m_oImports.end());

        for (auto& v : m_oImports)
        {
            ofs << std::endl << "\t\"" << v << "\"";
        }

        if (!m_oImports.empty()) ofs << std::endl;

        ofs << std::endl << "\t\"gitee.com/itfriday/protopack/deps/go/pp\"" << std::endl;

        ofs << ")" << std::endl;
    }

    // 常量定义
    void GoGenerator::GenerateConst(std::ofstream& ofs)
    {
        if (m_oConsts.empty()) return;

        std::vector<std::string> nlist;
        std::vector<std::string> vlist;
        uint32_t maxNameLen = 0;
        uint32_t maxDeclareLen = 0;
        uint32_t maxValueLen = 0;

        ofs << std::endl << "const (" << std::endl;

        // 算出名称的最大长度，对齐类型
        for (auto& v: m_oConsts)
        {
            if (v->GetName().length() > maxNameLen) maxNameLen = v->GetName().length();
        }

        // 算出const定义的最大长度，对齐=号
        for (auto& v: m_oConsts)
        {
            std::stringstream ss;

            ss << "\t" << v->GetName();
            for (size_t j = 0; j < maxNameLen - v->GetName().length(); j++) ss << " ";
            ss << " ";

            if (v->GetType() != "string")
            {
                ss << GetGoType(v->GetType());
            }

            std::string sval = ss.str();
            if (sval.length() > maxDeclareLen) maxDeclareLen = sval.length();
            nlist.push_back(sval);
        }

        // 算出const value的最大长度，对齐注释
        for (size_t i = 0; i < nlist.size(); i++)
        {
            std::stringstream ss;
            auto v = m_oConsts[i];

            ss << nlist[i];
            for (size_t j = 0; j < maxDeclareLen - nlist[i].length(); j++) ss << " ";
            ss << " = ";

            if (v->GetType() == "string")
            {
                ss << "\"" << v->GetValue() << "\"";
            }
            else
            {
                ss << v->GetValue();
            }

            std::string sval = ss.str();
            if (sval.length() > maxValueLen) maxValueLen = sval.length();
            vlist.push_back(sval);
        }

        // 对齐注释，输出到文件
        for (size_t i = 0; i < vlist.size(); i++)
        {
            auto v = m_oConsts[i];
            std::string sval = vlist[i];

            ofs << sval;
            if (!v->GetDesc().empty())
            {
                for (size_t j = 0; j < maxValueLen - sval.size(); j++) ofs << " ";
                ofs << " // " << v->GetDesc();
            }
            ofs << std::endl;
        }

        ofs << ")" << std::endl;
    }

    // 生成枚举定义
    void GoGenerator::GenerateEnum(std::ofstream& ofs)
    {
        for (auto& v : m_oEnums)
        {
            ofs << std::endl;

            if (!v->GetDesc().empty()) ofs << "// " << v->GetName() << " " << v->GetDesc() << std::endl;
            ofs << "type " << v->GetName() << " pp.Enum" << std::endl << std::endl;
            GenerateEnumValues(ofs, v->GetName(), v->GetFields());
        }
    }

    // 生成枚举值
    void GoGenerator::GenerateEnumValues(std::ofstream& ofs, const std::string& enumName,
        const std::vector<std::shared_ptr<FieldInfo> >& vals)
    {
        if (vals.empty()) return;

        std::vector<std::string> vlist;
        uint32_t maxNameLen = 0;
        uint32_t maxValueLen = 0;

        ofs << "const (" << std::endl;

        // 算出名称的最大长度，对齐类型
        for (auto& v: vals)
        {
            if (v->GetName().length() > maxNameLen) maxNameLen = v->GetName().length();
        }

        // 算出const定义的最大长度，对齐=号
        for (auto& v: vals)
        {
            std::stringstream ss;

            ss << "\t" << v->GetName();
            for (size_t j = 0; j < maxNameLen - v->GetName().length(); j++) ss << " ";
            ss << " " << enumName << " = "  << v->GetValue();

            std::string sval = ss.str();
            if (sval.length() > maxValueLen) maxValueLen = sval.length();
            vlist.push_back(sval);
        }

        // 对齐注释，输出到文件
        for (size_t i = 0; i < vlist.size(); i++)
        {
            auto v = vals[i];
            std::string sval = vlist[i];

            ofs << sval;
            if (!v->GetDesc().empty())
            {
                for (size_t j = 0; j < maxValueLen - sval.size(); j++) ofs << " ";
                ofs << " // " << v->GetDesc();
            }
            ofs << std::endl;
        }

        ofs << ")" << std::endl;
    }

    // 生成结构体定义
    void GoGenerator::GenerateStruct(std::ofstream& ofs)
    {
        for (auto& v : m_oStructs)
        {
            std::string className = v->GetName();
            bool bUnion = v->GetType() == STRUCT_TYPE_ONEOF;

            ofs << std::endl;

            if (!v->GetDesc().empty()) ofs << "// " << className << " " << v->GetDesc() << std::endl;
            ofs << "type " << className << " struct {" << std::endl;
            // 类成员定义
            GenerateStructAttributes(ofs, v->GetFields());
            ofs << "}" << std::endl;
            GenerateStructNewFunction(ofs, className, bUnion, v->GetFields());
            GenerateStructGetFunctions(ofs, className, v->GetFields());
            GenerateStructMarshal(ofs, className, v->GetFields());
            GenerateStructUnMarshal(ofs, className, v->GetFields());
        }
    }

    // 生成结构体成员定义
    void GoGenerator::GenerateStructAttributes(std::ofstream& ofs, const std::vector<std::shared_ptr<FieldInfo> >& vals)
    {
        if (vals.empty()) return;

        std::vector<std::string> nlist;
        std::vector<std::string> jlist;
        uint32_t maxNameLen = 0;
        uint32_t maxDeclareLen = 0;
        uint32_t maxJsonLen = 0;

        // 算出名称的最大长度，对齐类型
        for (auto& v: vals)
        {
            if (v->GetName().length() > maxNameLen) maxNameLen = v->GetName().length();
        }

        // 算出属性定义的最大长度，对齐``号
        for (auto& v: vals)
        {
            std::stringstream ss;

            ss << "\t" << ToPascalCase(v->GetName());
            for (size_t j = 0; j < maxNameLen - v->GetName().length(); j++) ss << " ";
            ss << " " << GetStructType(v->GetType(), v->GetSubType());

            std::string sval = ss.str();
            if (sval.length() > maxDeclareLen) maxDeclareLen = sval.length();
            nlist.push_back(sval);
        }

        // 算出json定义的最大长度，对齐注释
        for (size_t i = 0; i < nlist.size(); i++)
        {
            std::stringstream ss;
            auto v = vals[i];

            ss << nlist[i];
            for (size_t j = 0; j < maxDeclareLen - nlist[i].length(); j++) ss << " ";

            std::string tmp = v->GetJsonName().empty() ? v->GetName() : v->GetJsonName();
            if (!v->GetJsonOmitEmpty().empty()) tmp += "," + v->GetJsonOmitEmpty();

            ss << " `json:\"" << tmp << "\"`";

            std::string sval = ss.str();
            if (sval.length() > maxJsonLen) maxJsonLen = sval.length();
            jlist.push_back(sval);
        }

        // 对齐注释，输出到文件
        for (size_t i = 0; i < jlist.size(); i++)
        {
            auto v = vals[i];
            std::string sval = jlist[i];

            ofs << sval;
            if (!v->GetDesc().empty())
            {
                for (size_t j = 0; j < maxJsonLen - sval.size(); j++) ofs << " ";
                ofs << " // " << v->GetDesc();
            }
            ofs << std::endl;
        }
    }

    // 生成结构体New定义
    void GoGenerator::GenerateStructNewFunction(std::ofstream& ofs, const std::string& structName, bool bUnion,
        const std::vector<std::shared_ptr<FieldInfo> >& vals)
    {
        ofs << std::endl << "func New" << structName << "() *" << structName << "{" << std::endl
            << "\treturn &" << structName << "{";

        std::stringstream ss;
        uint32_t maxNameLen = 0;

        for (auto& v: vals)
        {
            if (v->GetType() == "array") continue;

            if ((!bUnion && m_oStructMap.find(v->GetType()) != m_oStructMap.end()) || !v->GetDefault().empty())
            {
                if (v->GetName().length() > maxNameLen) maxNameLen = v->GetName().length();
            }
        }

        for (auto& v: vals)
        {
            if (v->GetType() == "array") continue;

            if (!bUnion)
            {
                auto siter = m_oStructMap.find(v->GetType());
                if (siter != m_oStructMap.end())
                {
                    ss << "\t\t" << ToPascalCase(v->GetName()) << ": ";
                    for (size_t j = 0; j < maxNameLen - v->GetName().length(); j++) ss << " ";
                    ss << siter->second << "()," << std::endl;
                    continue;
                }
            }

            auto iter = TYPES_MAP.find(v->GetType());
            auto eiter = m_oEnumList.find(v->GetType());

            if (v->GetDefault().empty()) continue;

            ss << "\t\t" << ToPascalCase(v->GetName()) << ": ";

            for (size_t j = 0; j < maxNameLen - v->GetName().length(); j++) ss << " ";

            if (v->GetType() == "string")
            {
                ss << "\"" << v->GetDefault() << "\"," << std::endl;
            }
            else if (iter != TYPES_MAP.end())
            {
                ss << iter->second << "(" << v->GetDefault() << ")," << std::endl;
            }
            else if (eiter != m_oEnumList.end())
            {
                ss << *eiter << "(" << v->GetDefault() << ")," << std::endl;
            }
        }

        std::string def = ss.str();
        if (def.empty())
        {
            ofs << "}" << std::endl;
        }
        else
        {
            ofs << std::endl << def << "\t}" << std::endl;
        }

        ofs << "}" << std::endl;
    }

    // 生成结构体成员Get定义
    void GoGenerator::GenerateStructGetFunctions(std::ofstream& ofs, const std::string& structName,
        const std::vector<std::shared_ptr<FieldInfo> >& vals)
    {
        if (vals.empty()) return;

        for (auto& v: vals)
        {
            ofs << std::endl << "func (x *" << structName << ") Get" << ToPascalCase(v->GetName()) << "() " <<
                GetStructType(v->GetType(), v->GetSubType()) << " {" << std::endl;

            ofs << "\tif x != nil {" << std::endl << "\t\treturn x." << ToPascalCase(v->GetName()) << std::endl
                << "\t}" << std::endl;

            if (v->GetType() == "string")
            {
                ofs << "\treturn \"\"" << std::endl;
            }
            else
            {
                if (TYPES_MAP.find(v->GetType()) != TYPES_MAP.end() || m_oEnumList.find(v->GetType()) != m_oEnumList.end())
                {
                    ofs << "\treturn 0" << std::endl;
                }
                else
                {
                    ofs << "\treturn nil" << std::endl;
                }
            }

            ofs << "}" << std::endl;
        }
    }

    // 生成结构体成员Marshal定义
    void GoGenerator::GenerateStructMarshal(std::ofstream& ofs, const std::string& structName,
        const std::vector<std::shared_ptr<FieldInfo> >& vals)
    {
        ofs << std::endl << "func (x *" << structName << ") Marshal(buf *bytes.Buffer) error {"
            << std::endl << "\tif x == nil {" << std::endl << "\t\treturn nil"
            << std::endl << "\t}" << std::endl;

        for (auto& v: vals)
        {
            if (v->GetType() == "array")
            {
                std::string funcName = m_oEnumList.find(v->GetSubType()) != m_oEnumList.end() ? "ArrayEnumMarshal" : "ArrayFieldMarshal";
                ofs << "\tif err := pp." << funcName << "(buf, uint64(" << v->GetTag() << "), x."
                    << ToPascalCase(v->GetName()) << "); err != nil {" << std::endl
                    << "\t\treturn err" << std::endl << "\t}" << std::endl;
            }
            else if (TYPES_MAP.find(v->GetType()) != TYPES_MAP.end())
            {
                ofs << "\tif err := pp.FieldMarshal(buf, uint64(" << v->GetTag() << "), x."
                    << ToPascalCase(v->GetName()) << "); err != nil {" << std::endl
                    << "\t\treturn err" << std::endl << "\t}" << std::endl;
            }
            else if (m_oEnumList.find(v->GetType()) != m_oEnumList.end())
            {
                ofs << "\tif err := pp.FieldMarshal(buf, uint64(" << v->GetTag() << "), (*pp.Enum)(&x."
                    << ToPascalCase(v->GetName()) << ")); err != nil {" << std::endl
                    << "\t\treturn err" << std::endl << "\t}" << std::endl;
            }
            else
            {
                ofs << "\tif x." << ToPascalCase(v->GetName()) << " != nil {" << std::endl;
                ofs << "\t\tif err := pp.FieldMarshal(buf, uint64(" << v->GetTag() << "), x."
                    << ToPascalCase(v->GetName()) << "); err != nil {" << std::endl
                    << "\t\t\treturn err" << std::endl << "\t\t}" << std::endl;
                ofs << "\t}" << std::endl;
            }
        }

        ofs << "\treturn nil" << std::endl << "}" << std::endl;
    }

    // 生成结构体成员UnMarshal定义
    void GoGenerator::GenerateStructUnMarshal(std::ofstream& ofs, const std::string& structName,
        const std::vector<std::shared_ptr<FieldInfo> >& vals)
    {
        ofs << std::endl << "func (x *" << structName << ") UnMarshal(tag uint64, bytes []byte, start, end int) error {"
            << std::endl << "\tif x == nil {" << std::endl << "\t\treturn pp.ErrUnmarshal"
            << std::endl << "\t}" << std::endl;

        if (!vals.empty()) ofs << std::endl << "\tswitch tag {" << std::endl;

        for (auto& v: vals)
        {
            ofs << "\tcase uint64(" << v->GetTag() << "):" << std::endl;

            auto iter = UNMARSHAL_MAP.find(v->GetType());
            if (iter != UNMARSHAL_MAP.end())
            {
                ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = pp." << iter->second << "(bytes, start)" << std::endl;
            }
            else if (v->GetType() == "string")
            {
                ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = pp.StringFieldUnmarshal(bytes, start, end)" << std::endl;
            }
            else if (v->GetType() == "array")
            {
                auto siter = UNMARSHAL_MAP.find(v->GetSubType());
                auto stiter = m_oStructMap.find(v->GetSubType());
                if (siter != UNMARSHAL_MAP.end())
                {
                    ofs << "\t\tv := pp." << siter->second << "(bytes, start)" << std::endl;
                    ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = append(x." << ToPascalCase(v->GetName())
                        << ", v)" << std::endl;
                }
                else if (v->GetSubType() == "string")
                {
                    ofs << "\t\tv := pp.StringFieldUnmarshal(bytes, start, end)" << std::endl;
                    ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = append(x." << ToPascalCase(v->GetName())
                        << ", v)" << std::endl;
                }
                else if (m_oEnumList.find(v->GetSubType()) != m_oEnumList.end())
                {
                    ofs << "\t\tvar v " << v->GetSubType() << std::endl;
                    ofs << "\t\tif err := pp.FieldUnMarshal(bytes, start, end, (*Enum)(&v), false); err != nil {" << std::endl
                        << "\t\t\treturn err" << std::endl << "\t\t}" << std::endl;
                    ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = append(x." << ToPascalCase(v->GetName())
                        << ", v)" << std::endl;
                }
                else if (stiter != m_oStructMap.end())
                {
                    ofs << "\t\tv := " << stiter->second << "()" << std::endl;
                    ofs << "\t\tif err := pp.FieldUnMarshal(bytes, start, end, v, false); err != nil {" << std::endl
                        << "\t\t\treturn err" << std::endl << "\t\t}" << std::endl;
                    ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = append(x." << ToPascalCase(v->GetName())
                        << ", v)" << std::endl;
                }
                else
                {
                    ofs << "\t\tvar v " << v->GetSubType() << std::endl;
                    ofs << "\t\tif err := pp.FieldUnMarshal(bytes, start, end, &v, false); err != nil {" << std::endl
                        << "\t\t\treturn err" << std::endl << "\t\t}" << std::endl;
                    ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = append(x." << ToPascalCase(v->GetName())
                        << ", &v)" << std::endl;
                }
            }
            else if (m_oEnumList.find(v->GetType()) != m_oEnumList.end())
            {
                ofs << "\t\treturn (*pp.Enum)(&x." << ToPascalCase(v->GetName())
                    << ").UnMarshal(0, bytes, start, end)" << std::endl;
            }
            else
            {
                auto stiter = m_oStructMap.find(v->GetType());
                std::string newFunc = stiter == m_oStructMap.end() ? "new(" + v->GetType() + ")" : stiter->second + "()";
                ofs << "\t\tx." << ToPascalCase(v->GetName()) << " = " << newFunc << std::endl
                    << "\t\treturn pp.FieldUnMarshal(bytes, start, end, x."
                    << ToPascalCase(v->GetName()) << ", false)" << std::endl;
            }
        }

        if (!vals.empty()) ofs << "\tdefault:" << std::endl << "\t}" << std::endl << std::endl;

        ofs << "\treturn nil" << std::endl << "}" << std::endl;
    }

    // 返回类成员定义
    std::string GoGenerator::GetStructType(const std::string& ftype, const std::string& subtype)
    {
        if (ftype == "array")
        {
            auto iter = TYPES_MAP.find(subtype);
            if (iter != TYPES_MAP.end()) return "[]" + iter->second;
            auto eiter = m_oEnumList.find(subtype);
            if (eiter != m_oEnumList.end()) return "[]" + subtype;
            return "[]*" + GetImportValue(subtype);
        }

        auto iter = TYPES_MAP.find(ftype);
        if (iter != TYPES_MAP.end())
        {
            return iter->second;
        }

        auto eiter = m_oEnumList.find(ftype);
        if (eiter != m_oEnumList.end()) return ftype;

        return "*" + GetImportValue(ftype);
    }

    std::string GoGenerator::GetImportValue(const std::string& ftype)
    {
        return ftype;
    }
};
